/******************************************************************************

    Copyright (C) 2019 ErisedMedia
    All Rights Reserved.

    This source code and any compilation or derivative thereof is the
    proprietary information of ErisedMedia and is
    confidential in nature.
    Under no circumstances is this software to be exposed to or placed under an
    Open Source License of any type without the expressed written permission of
    ErisedMedia.

******************************************************************************/

/******************************************************************************

    File:   FaceRecDB.cc
    Brief:  FaceRecDB实现，人脸识别.

******************************************************************************/

/******************************************************************************

    Rev     Date        Author      Comments
--------------------------------------------------------------------------------
    001     20190526    RunyuanYe        Original

--------------------------------------------------------------------------------
    For consistency and standardisation retain the Section Separators.

*******************************************************************************/


/******************************************************************************
*
*   INCLUDE FILES
*
******************************************************************************/

#include "Head.h"
#include "FaceRecDB.h"
#include "Tools.h"

/******************************************************************************
*
*    LOCAL MACROS
*
******************************************************************************/


/******************************************************************************
*
*    LOCAL TYPEDEFS
*
******************************************************************************/


/******************************************************************************
*
*    STATIC DATA
*
******************************************************************************/


/******************************************************************************
*
*    EXPORTED DATA
*
******************************************************************************/


/******************************************************************************
*
*    EXPORTED FUNCTION
*
******************************************************************************/


/******************************************************************************
*
*    STATIC FUNCTION PROTOTYPES
*
******************************************************************************/


/******************************************************************************
*
*    FUNCTION IMPLEMENTATION
*
******************************************************************************/


FaceRecDB::FaceRecDB(float sameFaceScoreThreshold, int similarityTopK, int searchMode, int searchThreadnum)
{
    std::lock_guard<std::mutex> guard(m_dbLock);

    m_similarityTopK    = similarityTopK;
    m_sameFaceScoreThreshold = sameFaceScoreThreshold;
    m_context   = nullptr;
    m_callback  = nullptr;

    m_searchMode = searchMode;
    m_searchThreadnum = searchThreadnum;

    m_faceId = 0LL;
}

FaceRecDB::~FaceRecDB()
{
}

void FaceRecDB::RegisterFeatureCompare(FeatureCompareCB callback, void* context)
{
    std::lock_guard<std::mutex> guard(m_dbLock);

    m_context   = context;
    m_callback  = callback;
}

int FaceRecDB::InsertFace(std::shared_ptr<FaceFeatureInfo> feature, FaceRecDBParam* param, bool isRegister)
{
    if(isRegister)
    {
        feature->idx = 0LL;
        m_FaceIdFeaturesMap.push_back({feature});
        std::vector<std::shared_ptr<FaceFeatureInfo>> infos;
        m_MoveOutFaceIdFeaturesMap.push_back(infos);
    }
    else
    {
        auto& features = m_FaceIdFeaturesMap[feature->faceId];
        feature->idx = features.size();
        features.push_back(feature);
    }
    // m_FaceIdFeaturesMap[feature->faceId].push_back(feature);

    return 0;
}

int FaceRecDB::UpdateFace(std::shared_ptr<FaceFeatureInfo> feature, std::shared_ptr<FaceFeatureInfo> matchFeature, FaceRecDBParam* param)
{    
    // 更新缓存
    long long faceId = matchFeature->faceId;
    feature->idx = matchFeature->idx;
    m_FaceIdFeaturesMap[faceId][matchFeature->idx] = feature;
    m_MoveOutFaceIdFeaturesMap[faceId].push_back(matchFeature);

    return 0;
}

int FaceRecDB::MatchFace(std::shared_ptr<FaceFeatureInfo> feature, std::shared_ptr<FaceFeatureInfo> matchFeature, FaceRecDBParam* param)
{
    long long faceId = matchFeature->faceId;
    m_MoveOutFaceIdFeaturesMap[faceId].push_back(feature);

    return 0;
}

int FaceRecDB::SearchSimilarityTopK(void* feature, int featureSize, std::vector<std::pair<float, std::shared_ptr<FaceFeatureInfo>>>& features, int topk, float scoreThreshold)
{
    if(topk > 1)
    {
        std::vector<std::pair<float, int>> score_index_vec;
        std::vector<std::shared_ptr<FaceFeatureInfo>> feature_vec;
        int index = 0;
        for(auto& m : m_FaceIdFeaturesMap)
        {
            auto& features = m;

            for(auto info : features)
            {
                float score;
                
                score = m_callback(feature, (char*)(info->feature), featureSize, m_context);

                if(scoreThreshold <= score)
                {
                    score_index_vec.emplace_back(score, index++);
                    feature_vec.push_back(info);
                }
            }
        }

        int size = score_index_vec.size();
        if(topk > size)
        {
            topk = size;
        }
        if(size > 1)
        {
            std::stable_sort(score_index_vec.begin(), score_index_vec.end(),
                        [](const std::pair<float, int>& lhs, const std::pair<float, int>& rhs) {
                            return lhs.first > rhs.first;
                        });
        }

        for(int i = 0; i < topk; ++i)
        {
            std::pair<float, int>& score_idx = score_index_vec[i];
            features.emplace_back(score_idx.first, feature_vec[score_idx.second]);
        }
    }
    else
    {
        float max_score = 0.0f;
        std::shared_ptr<FaceFeatureInfo> max_score_info = nullptr;
        for(auto& m : m_FaceIdFeaturesMap)
        {
            auto& features = m;

            for(auto info : features)
            {
                float score;
                
                score = m_callback(feature, (char*)(info->feature), featureSize, m_context);

                if(scoreThreshold <= score)
                {
                    if(max_score <= score)
                    {
                        max_score = score;
                        max_score_info = info;
                    }
                }
            }            
        }
        if(max_score_info)
        {
            features.emplace_back(max_score, max_score_info);
        }
    }

    return 0;
}

int FaceRecDB::SearchSimilarityTopKMT(void* feature, int featureSize, std::vector<std::pair<float, std::shared_ptr<FaceFeatureInfo>>>& features, int topk, float scoreThreshold)
{
    mutex lock;
    int threadnum = m_searchThreadnum;
    int processCount = m_FaceIdFeaturesMap.size();

    if(processCount < 1)
    {
        return 0;
    }

    if (threadnum > processCount)
    {
        threadnum = processCount;
    }

    if(topk > 1)
    {
        std::vector<std::pair<float, int>> score_index_vec;
        std::vector<std::shared_ptr<FaceFeatureInfo>> feature_vec;
        int index = 0;

        std::thread workers[threadnum];
        for (int i = 0; i < threadnum; i++) {
            workers[i] = thread([&,i]() {
                int idx = i;

                std::vector<std::pair<float, std::shared_ptr<FaceFeatureInfo>>> _score_feature_vec;

                while (idx < processCount)
                {
                    auto& features = m_FaceIdFeaturesMap[idx];

                    for(auto info : features)
                    {
                        float score;
                        
                        score = m_callback(feature, (char*)(info->feature), featureSize, m_context);

                        if(scoreThreshold <= score)
                        {
                            _score_feature_vec.emplace_back(score, info);
                        }
                    }

                    idx += threadnum;
                }

                lock.lock();
                for (auto& p : _score_feature_vec)
                {
                    score_index_vec.emplace_back(p.first, index++);
                    feature_vec.push_back(p.second);
                }
                lock.unlock();
            });
        }

        for (auto &w : workers) {
            if (w.joinable()) w.join();
        }

        int size = score_index_vec.size();
        if(topk > size)
        {
            topk = size;
        }
        if(size > 1)
        {
            std::stable_sort(score_index_vec.begin(), score_index_vec.end(),
                        [](const std::pair<float, int>& lhs, const std::pair<float, int>& rhs) {
                            return lhs.first > rhs.first;
                        });
        }

        for(int i = 0; i < topk; ++i)
        {
            std::pair<float, int>& score_idx = score_index_vec[i];
            features.emplace_back(score_idx.first, feature_vec[score_idx.second]);
        }
    }
    else
    {
        float max_score = 0.0f;
        std::shared_ptr<FaceFeatureInfo> max_score_info = nullptr;

        std::thread workers[threadnum];
        for (int i = 0; i < threadnum; i++) {
            workers[i] = thread([&,i]() {
                int idx = i;

                float _max_score = 0.0f;
                std::shared_ptr<FaceFeatureInfo> _max_score_info = nullptr;

                while (idx < processCount)
                {
                    auto& features = m_FaceIdFeaturesMap[idx];

                    for(auto info : features)
                    {
                        float score;
                        
                        score = m_callback(feature, (char*)(info->feature), featureSize, m_context);

                        if(scoreThreshold <= score)
                        {
                            if(_max_score <= score)
                            {
                                _max_score = score;
                                _max_score_info = info;
                            }
                        }
                    }

                    idx += threadnum;
                }

                lock.lock();
                if(max_score < _max_score)
                {
                    max_score = _max_score;
                    max_score_info = _max_score_info;
                }
                lock.unlock();
            });
        }

        for (auto &w : workers) {
            if (w.joinable()) w.join();
        }

        if(max_score_info)
        {
            features.emplace_back(max_score, max_score_info);
        }
    }

    return 0;
}

int FaceRecDB::SearchFaceRecDB(FaceRecDBParam* param, FaceRecDBResult* result)
{
    std::lock_guard<std::mutex> guard(m_dbLock);
    int status = FaceRecStatus_Register;
    // float score= 1.0;

    std::shared_ptr<FaceFeatureInfo> matchFeature;
    std::vector<std::pair<float, std::shared_ptr<FaceFeatureInfo>>> topKFeatures;

    START_PROFILE_TIME_EX(pt0, "SearchSimilarityTopK", 100);
    if(0==m_searchMode)
    {
        if(0 != SearchSimilarityTopK(param->info->feature, param->info->featureSize, topKFeatures, m_similarityTopK, m_sameFaceScoreThreshold))
        {
            END_PROFILE_TIME(pt0);
            return -1;
        }
    }
    else
    {
        if(0 != SearchSimilarityTopKMT(param->info->feature, param->info->featureSize, topKFeatures, m_similarityTopK, m_sameFaceScoreThreshold))
        {
            END_PROFILE_TIME(pt0);
            return -1;
        }
    }
    
    END_PROFILE_TIME(pt0);
    if(topKFeatures.size() == 0)
    {
        status = FaceRecStatus_Register;
    }
    else
    {
        long long matchFaceId; // 最匹配的人脸ID
        if (m_similarityTopK > 1)
        {
            // 将topK的相同人脸ID的不同特征得分放一起
            std::map<long long, std::vector<float>> faceIdMap;
            for(auto& topKFeature : topKFeatures)
            {
                auto& faceId = topKFeature.second->faceId;
                faceIdMap[faceId].push_back(topKFeature.first);
                // if(faceIdMap.find(faceId) == faceIdMap.end())
                // {// 没找到
                //     std::vector<float> scores;
                //     scores.push_back(topKFeature.first);
                //     faceIdMap[faceId] = scores;
                // }
                // else
                // {// 找到
                //     faceIdMap[faceId].push_back(topKFeature.first);
                // }
            }

            int maxCount = 0;
            for(auto& fim: faceIdMap)
            {
                int count = fim.second.size();
                // TODO: 尝试改为有多个离最大的小于一定范围内的，看哪一个多归哪一个？
                if(count > maxCount)
                {// 选择匹配特征数量最多的人脸ID
                    maxCount = count;
                    matchFaceId = fim.first; 
                }
                else if(count == maxCount)
                {// 如果数量一样多，则选择特征之和最大的人脸ID
                    auto& fim1 = faceIdMap[matchFaceId];
                    float scoreSum1 = std::accumulate(fim1.begin(), fim1.end(), 0.0);
                    float scoreSum2 = std::accumulate(fim.second.begin(), fim.second.end(), 0.0);
                    if(scoreSum2 > scoreSum1)
                    {
                        maxCount = count;
                        matchFaceId = fim.first; 
                    }
                }
            }
        }
        else
        {
            matchFaceId = topKFeatures[0].second->faceId;
        }

        // 找出此匹配的人脸的所有人脸特征
        std::vector<std::shared_ptr<FaceFeatureInfo>>& matchFeatures = m_FaceIdFeaturesMap[matchFaceId];

        // 如果当前人脸的特征数量大于等于每个人脸最大特征数量需要替换其中一个
        if( MAX_FEATURE_COUNT_PRE_FACE <= matchFeatures.size() )
        {
            // 两两计算相似度，找出相似度之和最小的作为离群人脸
            std::vector<void*> features;
            features.push_back((char*)(param->info->feature));
            for(auto& mfeature: matchFeatures)
            {
                features.push_back((char*)(mfeature->feature));
            }
            int count = features.size();
            std::vector<std::vector<float>> scoreMat;
            scoreMat.assign(count, std::vector<float>(count));
            for(int i = 0; i < count; ++i)
            {
                scoreMat[i][i] = 0.0;
                for(int j = i+1; j < count; ++j)
                {
                    float score = m_callback(features[i], features[j], param->info->featureSize, m_context);
                    scoreMat[i][j] = score;
                    scoreMat[j][i] = score;
                }
            }
            float minScoreSum = std::accumulate(scoreMat[0].begin(), scoreMat[0].end(), 0.0);
            int matchIdx = 0;
            float maxScore = 0.0;
            int maxScoreId = 1;
            for(int i = 1; i < count; ++i)
            {
                float sum = std::accumulate(scoreMat[i].begin(), scoreMat[i].end(), 0.0);
                if(minScoreSum > sum)
                {
                    minScoreSum = sum;
                    matchIdx = i;
                }
                if(maxScore < scoreMat[0][i])
                {
                    maxScore = scoreMat[0][i];
                    maxScoreId = i;
                }
            }
            if(matchIdx == 0)
            {
                // matchFeature = matchFeatures[0]; // 随便选一个
                matchFeature = matchFeatures[maxScoreId-1]; // 选最匹配的
                status = FaceRecStatus_Match;
            }
            else
            {
                matchFeature = matchFeatures[matchIdx-1]; // 因为matchIdx是把新输入的特征也算在里面，所以需要减1
                status = FaceRecStatus_Update;
            }
            // std::cout << "Score: " << maxScore;
            // for (int i = 1; i < count; i++)
            // {
            //     std::cout << ", " << scoreMat[0][i];
            // }
            // std::cout << std::endl;
            // score = maxScore;
        }
        else
        {
            matchFeature = matchFeatures[0]; // 随便选一个即可
            status = FaceRecStatus_Add;
        }
    }

    std::shared_ptr<FaceFeatureInfo> feature = param->info;
    if(!feature) return -1;

    if(FaceRecStatus_Match == status)
    {
        feature->faceId = matchFeature->faceId;
        if(result)
        {
            result->uuid    = m_FaceIdUUIDMap[matchFeature->faceId];
            result->faceId  = matchFeature->faceId;
        }
        matchFeature->matchCount += 1;
        START_PROFILE_TIME_EX(pt1, "MatchFace", 100);
        if(0 != MatchFace(feature, matchFeature, param))
        {
            END_PROFILE_TIME(pt1);
            DBG_PRINT_EX(DBG_LEVEL_ERROR1, "Match Face Error\n");
            return -1;
        }
        END_PROFILE_TIME(pt1);
        #if ENABLE_FACE_RECOGNITION_TEST
        DBG_PRINT_EX(DBG_LEVEL_1, "FaceRecStatus_Match %lld: %lld\n", matchFeature->faceId, matchFeature->matchCount);
        #endif
    }
    else
    {
        START_PROFILE_TIME_EX(pt3, "Register Add Update Face", 100);
        if(FaceRecStatus_Register == status)
        {
            feature->faceId = m_faceId++;
            feature->matchCount = 1;
            string uuid;
            GenerateUUIDEx(uuid);
            m_FaceIdUUIDMap.push_back(uuid);
            // m_FaceIdUUIDMap[feature->faceId] = uuid;
            #if ENABLE_FACE_RECOGNITION_TEST
            DBG_PRINT_EX(DBG_LEVEL_1, "FaceRecStatus_Register %lld: %lld, %s\n", feature->faceId, feature->matchCount, uuid.c_str());
            #endif
            if(0 != InsertFace(feature, param))
            {
                DBG_PRINT_EX(DBG_LEVEL_ERROR1, "Register Face Error\n");
                END_PROFILE_TIME(pt3);
                return -1;
            }
            m_faceFeatureCount++;
            m_faceCount++;
        }
        else if(FaceRecStatus_Add == status)
        {
            feature->faceId = matchFeature->faceId;
            feature->matchCount = 1;
            #if ENABLE_FACE_RECOGNITION_TEST
            DBG_PRINT_EX(DBG_LEVEL_1, "FaceRecStatus_Add %lld: %lld\n", feature->faceId, feature->matchCount);
            #endif
            if(0 != InsertFace(feature, param, false))
            {
                DBG_PRINT_EX(DBG_LEVEL_ERROR1, "Add Face Error\n");
                END_PROFILE_TIME(pt3);
                return -1;
            }
            m_faceFeatureCount++;
        }
        else if(FaceRecStatus_Update == status)
        {
            feature->faceId = matchFeature->faceId;
            feature->matchCount = matchFeature->matchCount + 1;
            #if ENABLE_FACE_RECOGNITION_TEST
            DBG_PRINT_EX(DBG_LEVEL_1, "FaceRecStatus_Update %lld: %lld\n", feature->faceId, feature->matchCount);
            #endif
            if(0 != UpdateFace(feature, matchFeature, param))
            {
                DBG_PRINT_EX(DBG_LEVEL_ERROR1, "Update Face Error\n");
                END_PROFILE_TIME(pt3);
                return -1;
            }
        }
        END_PROFILE_TIME(pt3);
        if(result)
        {
            result->uuid    = m_FaceIdUUIDMap[feature->faceId];
            result->faceId  = feature->faceId;
        }
    }
    if(result)
    {
        result->status  = status;
    }
    // DBG_PRINT_EX(DBG_LEVEL_1, "FaceID: %lld\n", m_faceId);
    return 0;
}

int FaceRecDB::SaveFaceRecDB(const string& name, const string& faceImageListFileName, const string& faceFeatureFileName)
{
    int count = 0;
    std::ofstream listFile(faceImageListFileName, std::ios::app);
    std::ofstream featFile(faceFeatureFileName, std::ios::app|std::ios::binary);
    if(listFile.is_open() && featFile.is_open())
    {
        // std::vector<std::vector<std::shared_ptr<FaceFeatureInfo>>>  m_FaceIdFeaturesMap;
        // std::vector<std::vector<std::shared_ptr<FaceFeatureInfo>>>  m_MoveOutFaceIdFeaturesMap;
        // std::vector<std::string>  m_FaceIdUUIDMap;
        count = m_FaceIdUUIDMap.size();
        for (int i = 0; i < count; i++)
        {
            auto& uuid = m_FaceIdUUIDMap[i];
            auto& feats = m_FaceIdFeaturesMap[i];
            auto& movefeats = m_MoveOutFaceIdFeaturesMap[i];
            string name_uuid = name + " " + uuid;
            for(auto feat : feats)
            {
                listFile << name_uuid << " " << feat->path << " 0\n";
                featFile.write(feat->feature, feat->featureSize);
            }
            for(auto feat : movefeats)
            {
                listFile << name_uuid << " " << feat->path << " 1\n";
                featFile.write(feat->feature, feat->featureSize);
            }
            int total_feat_count = feats.size() + movefeats.size();
            if(total_feat_count > 1000)
            {
                std::cout << name_uuid << " " << total_feat_count << std::endl;
            }
        }
        // std::cout << name << " FaceID Count: " << m_faceCount << ", Face Feature Count: " << m_faceFeatureCount << std::endl;
        // listFile.flush();
        // featFile.flush();
    }
    
    return count;
}